{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = '0'\n",
    "from transformers import AutoTokenizer,AutoModelForCausalLM"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "The model is automatically converting to bf16 for faster inference. If you want to disable the automatic precision, please manually add bf16/fp16/fp32=True to \"AutoModelForCausalLM.from_pretrained\".\n",
      "Try importing flash-attention for faster inference...\n"
     ]
    },
    {
     "data": {
      "application/json": {
       "ascii": false,
       "bar_format": null,
       "colour": null,
       "elapsed": 0.007950305938720703,
       "initial": 0,
       "n": 0,
       "ncols": null,
       "nrows": null,
       "postfix": null,
       "prefix": "Loading checkpoint shards",
       "rate": null,
       "total": 2,
       "unit": "it",
       "unit_divisor": 1000,
       "unit_scale": false
      },
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "aafd87ed60f5437f9ef34d6bb9db2804",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/2 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "alg_name = \"LoRA\"\n",
    "model_path = '/share/huggingface/'\n",
    "model_name = \"Qwen-1_8B-Chat\"\n",
    "device = 'auto'\n",
    "base_model = AutoModelForCausalLM.from_pretrained(model_path+model_name, trust_remote_code=True, device_map=device)\n",
    "tok = AutoTokenizer.from_pretrained(model_path+model_name, eos_token='<|endoftext|>', pad_token='<|endoftext|>',unk_token='<|endoftext|>', trust_remote_code=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[14880, 107354, 107976, 99470, 100045, 16744, 9370, 33447, 104670, 5122, 99316, 99270, 58364, 99748, 17714, 102030, 1773, 107421, 99316, 99270, 58364, 99748, 3837]\n",
      "[104128, 100040, 102030, 112112, 1773]\n",
      "[14880, 107354, 107976, 99470, 100045, 16744, 9370, 33447, 104670, 5122, 99316, 99270, 58364, 99748, 17714, 102030, 1773, 107421, 99316, 99270, 58364, 99748, 3837, 104128, 100040, 102030, 112112, 1773]\n"
     ]
    }
   ],
   "source": [
    "tok1 = tok.encode(\"请填写下列古诗文的后一句：克己复礼为仁。一日克己复礼，\",)\n",
    "tok2 = tok.encode(\"天下归仁焉。\")\n",
    "tok3 = tok.encode(\"请填写下列古诗文的后一句：克己复礼为仁。一日克己复礼，天下归仁焉。\")\n",
    "print(tok1)\n",
    "print(tok2)\n",
    "print(tok3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[0.0829, 0.1372, 0.1358, 0.1256, 0.0671, 0.0814, 0.1029, 0.1019, 0.0946,\n",
       "         0.0705]], grad_fn=<SoftmaxBackward0>)"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "\n",
    "# 定义一个简单的模型\n",
    "class SimpleModel(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(SimpleModel, self).__init__()\n",
    "        self.linear1 = nn.Linear(20, 30)  # 一个简单的线性层\n",
    "        self.linear2 = nn.Linear(30, 10)  # 一个简单的线性层\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.linear1(x)\n",
    "        x = nn.functional.relu(x)\n",
    "        x = self.linear2(x)\n",
    "        x = nn.functional.softmax(x, dim=1)\n",
    "        return x\n",
    "\n",
    "# 实例化模型、损失函数和优化器\n",
    "model = SimpleModel()\n",
    "loss_fn = nn.CrossEntropyLoss()  # 多分类问题，使用CrossEntropyLoss\n",
    "optimizer = optim.SGD(model.parameters(), lr=0.01)\n",
    "\n",
    "# 假设输入数据和标签\n",
    "input_data = torch.randn(1, 20)  # 随机生成一个样本\n",
    "labels = torch.rand(1, 10)  # 真实标签\n",
    "\n",
    "# 前向传播\n",
    "outputs = model(input_data)\n",
    "outputs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "linear1.weight torch.Size([30, 20]) torch.Size([30, 20])\n",
      "linear1.bias torch.Size([30]) torch.Size([30])\n",
      "linear2.weight torch.Size([10, 30]) torch.Size([10, 30])\n",
      "linear2.bias torch.Size([10]) torch.Size([10])\n"
     ]
    }
   ],
   "source": [
    "# 计算损失\n",
    "loss = loss_fn(outputs, labels)\n",
    "# 清零梯度\n",
    "optimizer.zero_grad()\n",
    "# 反向传播\n",
    "loss.backward(retain_graph=True)\n",
    "# 更新参数\n",
    "# optimizer.step()\n",
    "\n",
    "# 查看梯度\n",
    "for k, v in model.named_parameters():\n",
    "    print(k, v.shape, v.grad.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(tensor(8.7134, grad_fn=<DivBackward1>),\n",
       " tensor([[0.0829, 0.1372, 0.1358, 0.1256, 0.0671, 0.0814, 0.1029, 0.1019, 0.0946,\n",
       "          0.0705]], grad_fn=<SoftmaxBackward0>))"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "loss, outputs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[True, True, True, True, True, True, True, True, True, True, True, True,\n",
       "         True, True, True, True, True, True, True, True],\n",
       "        [True, True, True, True, True, True, True, True, True, True, True, True,\n",
       "         True, True, True, True, True, True, True, True],\n",
       "        [True, True, True, True, True, True, True, True, True, True, True, True,\n",
       "         True, True, True, True, True, True, True, True],\n",
       "        [True, True, True, True, True, True, True, True, True, True, True, True,\n",
       "         True, True, True, True, True, True, True, True],\n",
       "        [True, True, True, True, True, True, True, True, True, True, True, True,\n",
       "         True, True, True, True, True, True, True, True],\n",
       "        [True, True, True, True, True, True, True, True, True, True, True, True,\n",
       "         True, True, True, True, True, True, True, True],\n",
       "        [True, True, True, True, True, True, True, True, True, True, True, True,\n",
       "         True, True, True, True, True, True, True, True],\n",
       "        [True, True, True, True, True, True, True, True, True, True, True, True,\n",
       "         True, True, True, True, True, True, True, True],\n",
       "        [True, True, True, True, True, True, True, True, True, True, True, True,\n",
       "         True, True, True, True, True, True, True, True],\n",
       "        [True, True, True, True, True, True, True, True, True, True, True, True,\n",
       "         True, True, True, True, True, True, True, True],\n",
       "        [True, True, True, True, True, True, True, True, True, True, True, True,\n",
       "         True, True, True, True, True, True, True, True],\n",
       "        [True, True, True, True, True, True, True, True, True, True, True, True,\n",
       "         True, True, True, True, True, True, True, True],\n",
       "        [True, True, True, True, True, True, True, True, True, True, True, True,\n",
       "         True, True, True, True, True, True, True, True],\n",
       "        [True, True, True, True, True, True, True, True, True, True, True, True,\n",
       "         True, True, True, True, True, True, True, True],\n",
       "        [True, True, True, True, True, True, True, True, True, True, True, True,\n",
       "         True, True, True, True, True, True, True, True],\n",
       "        [True, True, True, True, True, True, True, True, True, True, True, True,\n",
       "         True, True, True, True, True, True, True, True],\n",
       "        [True, True, True, True, True, True, True, True, True, True, True, True,\n",
       "         True, True, True, True, True, True, True, True],\n",
       "        [True, True, True, True, True, True, True, True, True, True, True, True,\n",
       "         True, True, True, True, True, True, True, True],\n",
       "        [True, True, True, True, True, True, True, True, True, True, True, True,\n",
       "         True, True, True, True, True, True, True, True],\n",
       "        [True, True, True, True, True, True, True, True, True, True, True, True,\n",
       "         True, True, True, True, True, True, True, True],\n",
       "        [True, True, True, True, True, True, True, True, True, True, True, True,\n",
       "         True, True, True, True, True, True, True, True],\n",
       "        [True, True, True, True, True, True, True, True, True, True, True, True,\n",
       "         True, True, True, True, True, True, True, True],\n",
       "        [True, True, True, True, True, True, True, True, True, True, True, True,\n",
       "         True, True, True, True, True, True, True, True],\n",
       "        [True, True, True, True, True, True, True, True, True, True, True, True,\n",
       "         True, True, True, True, True, True, True, True],\n",
       "        [True, True, True, True, True, True, True, True, True, True, True, True,\n",
       "         True, True, True, True, True, True, True, True],\n",
       "        [True, True, True, True, True, True, True, True, True, True, True, True,\n",
       "         True, True, True, True, True, True, True, True],\n",
       "        [True, True, True, True, True, True, True, True, True, True, True, True,\n",
       "         True, True, True, True, True, True, True, True],\n",
       "        [True, True, True, True, True, True, True, True, True, True, True, True,\n",
       "         True, True, True, True, True, True, True, True],\n",
       "        [True, True, True, True, True, True, True, True, True, True, True, True,\n",
       "         True, True, True, True, True, True, True, True],\n",
       "        [True, True, True, True, True, True, True, True, True, True, True, True,\n",
       "         True, True, True, True, True, True, True, True]])"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch.autograd.grad(loss, model.linear1.weight, retain_graph=True)[0] == model.linear1.weight.grad"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "ke2torch23cu121",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.19"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
